package com.liru.neo4j_demo.config;

import dev.langchain4j.community.rag.content.retriever.neo4j.Neo4jGraph;
import dev.langchain4j.community.rag.content.retriever.neo4j.Neo4jText2CypherRetriever;
import dev.langchain4j.community.store.embedding.neo4j.Neo4jEmbeddingStore;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import org.neo4j.driver.*;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AConfig {
    @Bean
    OpenAiChatModel openAiChatModel(){
        return OpenAiChatModel.builder()
                .apiKey("sk-5267f25a129740eda30ba197444de657") // 注意：在生产环境中，建议将API密钥存储在配置文件或环境变量中
                .modelName("qwen-plus") // 使用的模型名称
                .temperature(0.7)
                .logRequests(true)
                .logResponses(true)
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1") // 阿里云DashScope兼容接口地址
                .build();
    }

    @Bean
    EmbeddingModel embeddingModel(){
        return new AllMiniLmL6V2EmbeddingModel();
    }

    @Bean
    Neo4jEmbeddingStore embeddingStore(EmbeddingModel embeddingModel){
        return Neo4jEmbeddingStore.builder()
                .withBasicAuth("neo4j://192.168.6.63:7687", "neo4j", "Secret123")
                .dimension(embeddingModel.dimension())
                .label("CustomLabel")
                .build();
    }

    @Bean
    Neo4jText2CypherRetriever retriever(OpenAiChatModel openAiChatModel){
        Driver driver = GraphDatabase.driver("neo4j://192.168.6.63:7687", AuthTokens.basic("neo4j", "Secret123"));
        try (Neo4jGraph graph = Neo4jGraph.builder()
                .driver(driver)
                .build()){

/*            // 创建数据库会话并执行Cypher查询
            try (Session session = driver.session()) { // 创建会话
                // 执行Cypher语句创建节点和关系：一本书《Dune》和它的作者Frank Herbert
                session.run("CREATE (book:Book {title: 'Dune'})<-[:WROTE]-(author:Person {name: 'Frank Herbert'})");
            }
            graph.refreshSchema();*/
            // 构建Neo4j文本到Cypher查询的检索器
            Neo4jText2CypherRetriever retriever = Neo4jText2CypherRetriever.builder()
                    .graph(graph) // 设置图数据库实例
                    .chatModel(openAiChatModel) // 设置聊天模型
                    .build();
/*
            Query query = new Query("Who wrote Dune?");

            // 执行检索操作
            List<Content> contents = retriever.retrieve(query);

            System.out.println("Retrieved contents: " +  contents);*/

            return retriever;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Bean
    EmbeddingStoreIngestor embeddingStoreIngestor(EmbeddingModel embeddingModel, Neo4jEmbeddingStore embeddingStore){
        // 创建文档分割器
        DocumentSplitter documentSplitter = new DocumentByParagraphSplitter(1000, 0);

        // 创建嵌入存储注入器
        return EmbeddingStoreIngestor.builder()
                .documentSplitter(documentSplitter)
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();
    }
}
